from langchain.chains import create_sql_query_chain
from langchain.prompts import PromptTemplate
from boat_nl2sql.utiles import extract_sql_from_markdown, load_table_info_from_json
from boat_nl2sql.oracle_connection import get_db_connection, FewShotPromptTemplate
from boat_nl2sql.generate_echarts import visualize_sql_result, analyze_data_for_visualization, convert_sql_result_to_dataframe

# 在这里导入您的 LLM
from boat_nl2sql.llm_module import SiliconFlow

llm = SiliconFlow()

print("\n===== 测试 SQL 查询链 =====")

# 加载全部表信息
full_table_info = load_table_info_from_json("boat_nl2sql/oracle_table_info.json")
print(f"全部表信息中包含 {len(full_table_info)} 个表")

# 指定要包含的表
include_tables = [
    'BAYONET_BUSSINESS.BOAT_WARNING',
    'BAYONET_BUSSINESS.OFF_SITE_CASE',
    'BAYONET_BUSSINESS.ZD_WARNING_TYPE',
    'BAYONET_DYNAMIC.VIDEO_ANALYSIS',
    'BAYONET_DYNAMIC.DATAFUSION',
    'BAYONET_BASICS.SYS_BAYONET'
]

# 获取数据库连接
db = get_db_connection(
    sample_rows=2,
    custom_table_info=full_table_info,
    include_tables=include_tables
)

# 配置提示模板
SQL_PROMPT = PromptTemplate.from_template(
    """你是一个SQL专家。根据以下表结构：  
    {table_info}  

    问题：{input}  
    请严格按以下规则响应：  
    1. 只输出SQL代码  
    2. 使用标准Oracle SQL语法  
    3. 不要包含分号结尾  
    4. 最多返回{top_k}条记录  
    5. Oracle中使用ROWNUM来限制返回记录数  

    SQL查询："""
)
examples = [
    {
        "input": "查询名称中包含'船闸'的所有卡口基本信息，按名称排序。",
        "query": "SELECT ID, NAME, UNIT_ID, DES, PHOTOHTTP  FROM BAYONET_BASICS.SYS_BAYONET  WHERE NAME LIKE '%船闸%'  ORDER BY NAME",
    },
    {
        "input": "统计一月份各卡口点的船舶通过数量，按照通过船舶数量降序排列，只显示通过量前10的卡口。",
        "query": "SELECT * FROM (SELECT b.NAME AS 卡口名称, COUNT(d.CODE) AS 通过船舶数量  FROM BAYONET_BASICS.SYS_BAYONET b  JOIN BAYONET_DYNAMIC.DATAFUSION d ON b.ID = d.BAYONET_ID  WHERE d.PASSTIME >= TRUNC(SYSDATE, 'YYYY') AND d.PASSTIME < ADD_MONTHS(TRUNC(SYSDATE, 'YYYY'), 1) GROUP BY b.NAME  ORDER BY 通过船舶数量 DESC) WHERE ROWNUM <= 10",
    },
    {
        "input": "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。",
        "query": "WITH warning_counts AS (  SELECT t.NAME AS 预警类型,   COUNT(w.CODE) AS 预警次数  FROM BAYONET_BUSSINESS.ZD_WARNING_TYPE t  JOIN BAYONET_BUSSINESS.BOAT_WARNING w ON t.ID = w.WARNING_TYPE  GROUP BY t.NAME  )  SELECT    预警类型,预警次数, ROUND(预警次数 * 100 / SUM(预警次数) OVER(), 2) || '%' AS 占比 FROM warning_counts ORDER BY 预警次数 DESC "
    },
    {
        "input": "分析一下2025年1月油墩港(东大盈船闸)的流量组成。",
        "query": """SELECT 
    COUNT(da.CODE) AS TOTAL_FLOW,          -- 总流量
    SUM(CASE WHEN da.DIRECTION = 1 THEN 1 ELSE 0 END) AS UPSTREAM_FLOW, -- 上行流量
    SUM(CASE WHEN da.DIRECTION = 2 THEN 1 ELSE 0 END) AS DOWNSTREAM_FLOW -- 下行流量
FROM 
    BAYONET_DYNAMIC.DATAFUSION da
LEFT JOIN BAYONET_BASICS.SYS_BAYONET sb
    ON da.BAYONET_ID = sb.ID
WHERE 
    sb.NAME = '油墩港(东大盈船闸)' 
    AND da.PASSTIME >= TO_DATE('2025-01-01 00:00:00', 'YYYY-MM-DD HH24:MI:SS')
    AND da.PASSTIME <= TO_DATE('2025-01-31 23:59:59', 'YYYY-MM-DD HH24:MI:SS');"""
    },
    {
        "input": "查询油墩港(东大盈船闸)2025年 1月份各类预警的案件总数与占比",
        "query": """WITH filtered_data AS (
  SELECT 
    zwt.NAME AS WARNING_TYPE_NAME
  FROM 
    BAYONET_DYNAMIC.DATAFUSION df
  JOIN 
    BAYONET_BUSSINESS.BOAT_WARNING bw 
    ON df.CODE = bw.EVENTCODE
  JOIN 
    BAYONET_BUSSINESS.ZD_WARNING_TYPE zwt  
    ON bw.WARNING_TYPE = zwt.ID  
  JOIN 
    BAYONET_BASICS.SYS_BAYONET sb  -- 关联卡口基础信息表
    ON df.BAYONET_ID = sb.ID       -- 通过ID关联
  WHERE 
    sb.NAME = '油墩港(东大盈船闸)'  -- 直接使用卡口名称过滤
    AND df.PASSTIME >= DATE '2025-01-01' 
    AND df.PASSTIME < DATE '2025-02-01'
)

SELECT  WARNING_TYPE_NAME,  COUNT(*) AS record_count,  ROUND(COUNT(*) * 100.0 / SUM(COUNT(*)) OVER (),2) AS percentage FROM filtered_data GROUP BY WARNING_TYPE_NAME ORDER BY  WARNING_TYPE_NAME;
"""
    }

]

few_shot_example_prompt = PromptTemplate.from_template("输入问题: {input}\n生成SQL: {query}")
few_shot_prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=few_shot_example_prompt,
    prefix="你是一个Oracle SQL专家。请根据以下表结构生成对应问题的Oracle SQL语句，Oracle版本较低(11g及以下)，且尽量符合标准语法,不考虑ISDELETE和ISFINISH字段，如果没有强调年份，默认为2025年。\n\n以下是表结构（table_info）：\n{table_info}\n\n这是一些问题及其对应的SQL样例:",
    suffix="\n---\n问题: {input}\n对应SQL: ",
    input_variables=["input", "table_info"],
)

# 创建 SQL 查询链
generate_query = create_sql_query_chain(llm=llm, db=db, prompt=SQL_PROMPT, k=5)

# 示例查询
# question = "查询所有卡口的名称"
question = "统计一月份各卡口点的船舶通过数量，按照通过船舶数量降序排列，只显示通过量前10的卡口。"
# question = "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。"
# question = "查询名为'苏无锡货08168'的船舶最近一次通过系统的完整记录，包括通过卡口信息、预警情况和消息发送记录。"
# question = "分析一下2025年1月油墩港(东大盈船闸)的流量组成。"
# question = "查询油墩港(东大盈船闸)2025年 1月份各类预警的案件总数与占比"
print(f"问题: {question}")

few_shot_sql = few_shot_prompt.format(input=question, table_info=db.table_info)

llm_response = llm.invoke(few_shot_sql)
print("FewShot 生成的 SQL:", llm_response)

if "```sql" in llm_response:
    extracted_sql = extract_sql_from_markdown(llm_response)
    print("提取出的 SQL:")
    print(extracted_sql)
else:
    extracted_sql = llm_response.replace(";", "")

try:
    # 执行SQL
    sql_result = db.run(extracted_sql)
    print("SQL结果:", sql_result)

    if len(sql_result) > 1:
        # 可视化SQL结果
        output_file = visualize_sql_result(
            sql_result=sql_result,
            query=extracted_sql,
            question=question
        )

        if output_file:
            print(f"可视化已完成，请打开以下文件查看: {output_file}")

            # 生成分析结果
            TEXT_PROMPT = PromptTemplate.from_template(
                """根据问题{test_question}，SQL查询结果{sql_result}，以及生成的图表类型{chart_type}，  
                请提供以下内容：  
                1. 对数据的简明分析  
                2. 从数据中得出的关键见解  
                3. 为什么选择了这种图表类型来展示数据  
                4. 根据数据可能的建议或行动计划  

                请使用简洁专业的语言回答。"""
            )

            # 从可视化推荐中获取图表类型
            recommendation = analyze_data_for_visualization(
                convert_sql_result_to_dataframe(sql_result),
                extracted_sql,
                question
            )

            chart_type = recommendation["chart_type"]

            # 格式化PromptTemplate为字符串
            formatted_prompt = TEXT_PROMPT.format(
                test_question=question,
                sql_result=sql_result,
                chart_type=chart_type
            )

            analysis_result = llm.invoke(formatted_prompt)
            print("\n数据分析结果:")
            print(analysis_result)
    else:
        # TEXT_PROMPT = PromptTemplate.from_template(
        #     """根据问题{test_question}，SQL查询结果{sql_result}，
        #     请提供以下内容：
        #     1. 对数据的简明分析
        #     2. 从数据中得出的关键见解
        #     3. 根据数据可能的建议或行动计划
        #
        #     请使用简洁专业的语言回答。"""
        # )
        TEXT_PROMPT = PromptTemplate.from_template(
            """
    你是一名经验丰富的港航专业人士。请你根据问题 {test_question} 和提供的【结构化结果】{sql_result} ，以业务视角进行如下专业解读：

    若结果非空，请按以下总结：
    1. 数据的简明分析
    2. 关键业务见解

    若结果为空，请严格遵循如下标准作答：
    1. 数据的简明分析
    2. 分析可能的原因（数据缺失、查询日期条件限制等）  
    3. 提出2-3个相关的跟进问题，帮助进一步明确用户需求

    # 说明
     - 务必结构简介清晰、表达专业、直奔业务价值核心
     - 避免无效的数据复述
     - 避免输出Markdown格式
     - 避免出现表名称或表字段的英文名称
            """
        )

        # 格式化PromptTemplate为字符串
        formatted_prompt = TEXT_PROMPT.format(
            test_question=question,
            sql_result=sql_result,
        )

        analysis_result = llm.invoke(formatted_prompt)
        print("\n数据分析结果:")
        print(analysis_result)

    # TEXT_PROMPT = PromptTemplate.from_template(
    #     """根据问题{test_question}，和查询数据库返回的结果{sql_result}，用简介的语言，回答提出的问题。""")
    # # 格式化 PromptTemplate 为字符串
    # formatted_prompt = TEXT_PROMPT.format(test_question=question, sql_result=sql_result)
    #
    # result = llm.siliconflow_completions(model="deepseek-ai/DeepSeek-V3", prompt=formatted_prompt)
    # print("结果:", result)

except Exception as e:
    print(f"执行查询时出错: {e}")